package com.xiashitech.agent.instrumentation.httprequest;

import cn.hutool.core.util.StrUtil;
import com.xiashitech.agent.config.StartConfig;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.CharArrayWriter;
import java.io.IOException;
import java.io.PrintWriter;

public class XSHttpServletResponseWrapper extends HttpServletResponseWrapper {
    private ByteArrayOutputStream byteArrayOutputStream;
    private CharArrayWriter charArrayWriter;

    public XSHttpServletResponseWrapper(HttpServletResponse response) {
        super(response);
        byteArrayOutputStream = new ByteArrayOutputStream();
        charArrayWriter = new CharArrayWriter();
    }

    public byte[] getByteArray() {
        return byteArrayOutputStream.toByteArray();
    }

    public char[] getCharArray() {
        return charArrayWriter.toCharArray();
    }

    public void setByteArrayOutputStream(ByteArrayOutputStream byteArrayOutputStream) {
        this.byteArrayOutputStream = byteArrayOutputStream;
    }

    public void setCharArrayWriter(CharArrayWriter charArrayWriter) {
        this.charArrayWriter = charArrayWriter;
    }

    public void transferDataToResponse() {
        try {
            if(byteArrayOutputStream.toByteArray().length > 0) {
                super.getResponse().getOutputStream().write(byteArrayOutputStream.toByteArray());
            } else if(charArrayWriter.toCharArray().length > 0) {
                super.getResponse().getWriter().write(charArrayWriter.toCharArray());
            }
        } catch (Throwable e) {
            StartConfig.throwableStackPrinter(e);
        }
    }

    @Override
    public PrintWriter getWriter() {
        return new PrintWriter(charArrayWriter);
    }
    @Override
    public ServletOutputStream getOutputStream() {
        return new XSServletOutputStream(byteArrayOutputStream);
    }
    @Override
    public void finalize() throws Throwable {
        super.finalize();
        byteArrayOutputStream.close();
        charArrayWriter.close();
    }

    public String getJsonBody() {
        String jsonBody = "";
        if (StrUtil.containsIgnoreCase(getContentType(), "json")){
            if (getByteArray().length > 0) {
                jsonBody = new String(getByteArray());
            } else if (getCharArray().length > 0) {
                jsonBody = new String(getCharArray());
            }
        }
        return jsonBody;
    }
}
